from tensorflow.keras import datasets, layers, models
import tensorflow as tf
import client
import server
import numpy as np

# 定义一个卷积神经网络（CNN）模型，用于训练和测试
def CNN_model_factory():
  model = models.Sequential()
  # 第一层卷积层，30 个卷积核，卷积核大小为 3x3，激活函数为 relu
  model.add(layers.Conv2D(30, (3, 3), activation='relu', input_shape=(28, 28, 1)))
  # 第一层池化层，池化窗口大小为 2x2
  model.add(layers.MaxPooling2D((2, 2)))
  # 第二层卷积层，50 个卷积核，卷积核大小为 3x3，激活函数为 relu
  model.add(layers.Conv2D(50, (3, 3), activation='relu'))
  # 第二层池化层，池化窗口大小为 2x2
  model.add(layers.MaxPooling2D((2, 2)))
  # 将多维数据展平为一维向量
  model.add(layers.Flatten())
  # 全连接层，100 个神经元，激活函数为 relu
  model.add(layers.Dense(100, activation='relu'))
  # 输出层，10 个神经元，激活函数为 softmax，用于输出分类概率
  model.add(layers.Dense(10,activation='softmax'))
  return model


#正常运行联邦学习。
def run_no_attacks(root_data, root_label, clients_data, clients_label, nclient, x_test, y_test):
  # 学习率相关超参数
  beta  = 5e-2
  alpha = 5e-2/beta
  # 模型工厂函数，用于创建 CNN 模型
  model_factory = CNN_model_factory
  # 存储客户端对象
  clients = []
  # 创建客户端对象
  for i in range(0,nclient):
    clients.append(client.Client(model_factory, clients_data[i], clients_label[i], learning_rate=beta, R=1, batch_size=32))
  # 创建根服务器对象
  root_client = client.Client(model_factory, root_data, root_label, learning_rate=beta, R=1, batch_size=32)
  # server初始化
  server1 = server.Server(model_factory, select_client=nclient, iteration=200, alpha=alpha, beta=beta)
  # 调用server的 train 方法进行训练
  server1.train(clients,root_client,'no_attacks',x_test, y_test)
  return server1

#仅用根数据集训练（基准对比）。
def run_only_server(root_data, root_label, x_test, y_test):
  # 定义学习率 beta
  beta  = 5e-2
  # 计算超参数 alpha
  alpha = 5e-2/beta
  # 获取模型工厂函数
  model_factory = CNN_model_factory
  # 创建服务器对象
  server1 = server.Server(model_factory, select_client=100, iteration=200, alpha=alpha, beta=beta)
  # 创建根客户端对象
  root_client = client.Client(model_factory, root_data, root_label, learning_rate=beta, R=1, batch_size=32)
  # 调用server的 train_self 方法仅在根数据集上进行训练
  server1.train_self(root_client,'train_on_server', x_test, y_test)
  return server1


def run_smarter_adaptive_attack(root_data, root_label, clients_data, clients_label, nclient,
                                x_test, y_test,
                                # 新增参数以控制新的攻击策略
                                initial_benign_count_param,
                                malicious_config_param,
                                benign_effectiveness_threshold_param,
                                cos_candidates_param,
                                expr_id_suffix="smarter_adaptive_attack"):  # 用于日志的标识

  beta = 5e-2
  alpha = 5e-2 / beta
  model_factory = CNN_model_factory
  clients_list = []
  for i in range(nclient):
    clients_list.append(
      client.Client(model_factory, clients_data[i], clients_label[i], learning_rate=beta, R=1, batch_size=32))

  root_client_obj = client.Client(model_factory, root_data, root_label, learning_rate=beta, R=1, batch_size=32)

  server_instance = server.Server(model_factory, select_client=nclient, iteration=200, alpha=alpha,
                                  beta=beta)  # iteration可配置

  # 构建实验名称
  exp_name = f'FLTrust_{expr_id_suffix}'
  # 可以考虑将更多配置参数加入exp_name中，如果需要区分不同配置的实验结果

  print(f"\n--- 开始运行 FLTrust ({exp_name}) ---")

  # 调用服务器端新的训练方法
  server_instance.train_smarter_adaptive_attack(  # 确保方法名与server.py中定义的一致
    clients=clients_list,
    root_client=root_client_obj,
    expr_basename=exp_name,
    x_test=x_test,
    y_test=y_test,
    initial_benign_count=initial_benign_count_param,
    malicious_config=malicious_config_param,
    benign_effectiveness_threshold=benign_effectiveness_threshold_param,
    cos_candidates=cos_candidates_param
  )
  print(f"--- 完成 FLTrust ({exp_name}) 的训练 ---")
  return server_instance